# TODO 1.将查询mongodb的代码按业务进行封装
# TODO 2.是否将crawler代码按不同数据源封装到此处

import pandas as pd
from AStock.ASSetting import settings
from AStock.ASUtil import date_to_timestamp, date_to_bar_format
from AStock.ASData import (
    financial_dict_english_to_id,
    financial_dict_id_to_chinese,
    financial_dict_id_to_english
)
from datetime import datetime, timedelta


def ASQuery_stock_info(
        code=None,
        market=None,
        fields=['code', 'name', 'industry', 'area'],
        sorting_field='code'
):
    coll = settings.database.stock_info_tushare
    filter = {}
    projection = {'_id': 0}
    if code and isinstance(code, list):
        filter.update({'code': {'$in': code}})
    elif code and isinstance(code, str):
        filter.update({'code': code})
    if market and isinstance(market, str):
        # filter.update({'market': market})
        raise NotImplementedError('market')
    if fields:
        for f in fields:
            projection.__setitem__(f, 1)
    if sorting_field:
        res = coll.find(filter, projection).sort(sorting_field)
    else:
        res = coll.find(filter, projection)
    res = [i for i in res]
    df = pd.DataFrame(res)
    # if sorting_field:
    #    df = df.sort_values(sorting_field, ascending=True)

    return df


def ASQuery_research_report(
        code,
        begin_date,
        end_date,
        fields=None,
        limit=None
):
    coll = settings.database.research_report
    filter = {
        'stockCode': code,
        'publishDate': {
            '$gte': int(begin_date),
            '$lte': int(end_date)
        }
    }
    projection = {'_id': 0}
    if fields:
        for f in fields:
            projection.__setitem__(f, 1)
    if limit:
        # 降序排序，取limit个
        res = coll.find(filter, projection).sort('publishDate', -1).limit(limit)
    else:
        res = coll.find(filter, projection).sort('publishDate', -1)
    res = [i for i in res]
    df = pd.DataFrame(res)

    return df


def ASQuery_stock_adj(code, start, end):
    coll = settings.database.stock_adj

    if isinstance(code, str):
        code = [code]

    projection = {"_id": 0}

    cursor = coll.find(
        {
            'code': {'$in': code},
            "date": {
                "$lte": date_to_bar_format(end),
                "$gte": date_to_bar_format(start)
            }
        },
        projection,
        batch_size=10000
    )

    res = pd.DataFrame([item for item in cursor])

    return res


def ASQuery_stock_xdxr(
        code,
        fields=None
):
    '获取股票除权信息/数据库'
    if isinstance(code, str):
        code = [code]
    projection = {"_id": 0}
    if fields:
        for f in fields:
            projection.__setitem__(f, 1)
        if 'code' not in fields:
            projection.__setitem__('code', 1)
        if 'date' not in fields:
            projection.__setitem__('date', 1)
    collections = settings.database.stock_xdxr
    try:
        data = pd.DataFrame(
            [item for item in collections.find({'code': {"$in": code}
                                                }, projection)])
        data['date'] = pd.to_datetime(data['date'])

        df = data.drop_duplicates(
            ['date', 'code'],
            keep='last'
        ).set_index(['date',
                     'code'],
                    drop=False)
        if fields and 'code' not in fields:
            df = df.drop('code', axis=1)
        if fields and 'date' not in fields:
            df = df.drop('date', axis=1)
        return df
    except:
        return pd.DataFrame(
            data=[],
            columns=fields if fields else [
                'category',
                'category_meaning',
                'code',
                'date',
                'fenhong',
                'fenshu',
                'liquidity_after',
                'liquidity_before',
                'name',
                'peigu',
                'peigujia',
                'shares_after',
                'shares_before',
                'songzhuangu',
                'suogu',
                'xingquanjia'
            ]
        )


# 市值用除权除息计算，价格用前复权
def _fill_mv(df, start, end):
    df['date'] = pd.to_datetime(df['date'])
    if 'total_mv' not in df.columns:
        df['total_mv'] = None
    if 'circ_mv' not in df.columns:
        df['circ_mv'] = None
    df_mv = df[pd.isna(df['total_mv']) | pd.isna(df['circ_mv'])]
    codes = df_mv['code'].tolist()
    if not codes:
        return df
    # 计算前复权价格
    df_adj = ASQuery_stock_adj(codes, start, end)
    if not df_adj.empty:
        coll = settings.database.stock_day
        cursor = coll.find(
            {
                'code': {'$in': codes},
                "date_stamp": {
                    "$lte": date_to_timestamp(end),
                    "$gte": date_to_timestamp(start)
                }
            },
            {'close': 1, 'date': 1, 'code': 1},
            batch_size=10000
        )
        df_close = pd.DataFrame([item for item in cursor])
        df_close['date'] = pd.to_datetime(df_close['date'])
        df_mv = df_mv.drop('close', axis=1)
        df_mv = df_mv.join(df_close.set_index(['code', 'date']), on=['code', 'date'], how='left')

        df_adj['date'] = pd.to_datetime(df_adj['date'])
        df_mv = df_mv.join(df_adj.set_index(['code', 'date']), on=['code', 'date'], how='left')
        df_mv['close'] = df_mv['close'] * df_mv['adj']
        df_mv = df_mv.drop('adj', axis=1)
    # 使用除权除息数据计算市值
    df_mv = df_mv.set_index(['date', 'code'], drop=False)
    xdxr = ASQuery_stock_xdxr(codes, fields=['date', 'code', 'shares_after', 'liquidity_after', 'category'])
    mv = xdxr.query('category!=6').loc[:, ['shares_after', 'liquidity_after']].dropna()
    df_mv = pd.concat([df_mv, mv], axis=1)
    df_mv = df_mv.assign(
        shares=df_mv.shares_after.groupby(level=1).fillna(method='ffill'),
        liquidity=df_mv.liquidity_after.groupby(level=1).fillna(method='ffill')
    )
    df_mv['total_mv'] = df_mv.close * df_mv.shares * 10000
    df_mv['circ_mv'] = df_mv.close * df_mv.liquidity * 10000
    df_mv = df_mv.dropna(subset=['code'])
    df_mv = df_mv[['total_mv', 'circ_mv']]
    df = df.set_index(['date', 'code'], drop=False)
    df = df.combine_first(df_mv)
    df = df.reset_index(drop=True)
    """
    # 日线数据中可能没有总市值，这种情况取上季度财务报表064实收资本（或股本）乘以股价即为总市值
    codes = []
    for tup in df_notice_date_price.itertuples():
        if 'calcDateTotalMV' not in tup or not tup.calcDateTotalMV:
            codes.append(tup.code)
    if self._report_date[4:6] == '03':
        last_quarter = str(int(self._report_date[:4]) - 1) + '1231'
    elif self._report_date[4:6] == '06':
        last_quarter = self._report_date[:4] + '0331'
    elif self._report_date[4:6] == '09':
        last_quarter = self._report_date[:4] + '0630'
    elif self._report_date[4:6] == '12':
        last_quarter = self._report_date[:4] + '0930'
    else:
        assert 0
    filter_ = {'report_date': int(last_quarter), 'code': {'$in': codes}}

    coll = settings.database.financial
    try:
        cursor = coll.find(filter_, {'_id': 0, 'code': 1, '064': 1})
    except Exception as e:
        print(e)
        return None
    data = [item for item in cursor]
    df_mv = pd.DataFrame(data)
    df_mv = df_mv.rename(columns={'064': 'calcDateTotalMV'})
    df = df.set_index('code').combine_first(df_mv.set_index('code'))
    df['calcDateTotalMV'] = df['calcDateTotalMV'] * df['noticeDatePrice']
    return df_notice_date_price.set_index('code').combine_first(df)
    """
    return df


def ASQuery_stock_day(
        code,
        start,
        end,
        fields=None,
        type='bfq'
):
    """
    查询股票日价格信息
    :param code:
    :param start:
    :param end:
    :param fields: 要返回的列，注意 total_mv 和 circ_mv
                    在 stock_day 表中不存在的话，则使用复权数据进行计算
    :param type: bfq 不复权，  qfq 前复权
    :return: DataFrame
    """
    coll = settings.database.stock_day

    if isinstance(code, str):
        code = [code]

    projection = {"_id": 0}
    if fields:
        for f in fields:
            projection.__setitem__(f, 1)

        if 'code' not in fields:
            projection.__setitem__('code', 1)
        if 'date' not in fields:
            projection.__setitem__('date', 1)

    cursor = coll.find(
        {
            'code': {'$in': code},
            "date_stamp": {
                "$lte": date_to_timestamp(end),
                "$gte": date_to_timestamp(start)
            }
        },
        projection,
        batch_size=10000
    )

    res = pd.DataFrame([item for item in cursor])

    if type == 'qfq':
        df_adj = ASQuery_stock_adj(code, start, end)
        if not df_adj.empty:
            res = res.join(df_adj.set_index(['code', 'date']), on=['code', 'date'], how='left')
            for col in ['open', 'high', 'low', 'close']:
                if col in res.columns:
                    res[col] = res[col] * res['adj']
            res = res.drop('adj', axis=1)

    # total_mv 和 circ_mv 在 stock_day 表中不存在的话，则使用复权数据进行计算
    if not fields:
        res = _fill_mv(res, start, end)
    elif 'total_mv' in fields or 'circ_mv' in fields:
        res = _fill_mv(res, start, end)
        if 'total_mv' not in fields:
            res = res.drop('total_mv', axis=1)
        if 'circ_mv' not in fields:
            res = res.drop('circ_mv', axis=1)

    if fields:
        if 'code' not in fields:
            res = res.drop('code', axis=1)
        if 'date' not in fields:
            res = res.drop('date', axis=1)

    return res


def ASQuery_valuation(
        code=None,
        fields=None
):
    coll = settings.database.stock_valuation
    if code and isinstance(code, str):
        filter = {'code': code}
    elif code and isinstance(code, list):
        filter = {'code': {'$in': code}}
    else:
        filter = {}

    projection = {'_id': 0}
    if fields:
        for f in fields:
            projection.__setitem__(f, 1)
        if 'code' not in fields:
            projection.__setitem__('code', 1)

    cursor = coll.find(filter, projection).sort('code')

    res = [i for i in cursor]
    df = pd.DataFrame(res)

    return df


def ASQuery_trade_cal(date):
    coll = settings.database.trade_cal
    x = coll.find_one({'cal_date': int(date)})
    if x: return x['is_open'] == 1
    raise KeyError(date)


def ASQuery_nearest_trade_date(date, future=False, max_days=30):
    """
    查询离输入日期最近的交易日
    :param date: str 参考的日期  yyyymmdd
    :param future: True - 未来  False - 过去
    :param max_days: 最多查找的天数，防止死循环
    :return: yyyymmdd 或者 None
    """
    i = 0
    while i < max_days:
        if not ASQuery_trade_cal(date):
            cdate = datetime.strptime(date, '%Y%m%d')
            if future:
                cdate += timedelta(days=1)
            else:
                cdate -= timedelta(days=1)
            date = cdate.strftime('%Y%m%d')

        else:
            if not future and date == datetime.now().strftime('%Y%m%d') and datetime.now().hour < 17:
                # 直接取前一天，防止当天没有收盘
                cdate = datetime.strptime(date, '%Y%m%d')
                cdate -= timedelta(days=1)
                date = cdate.strftime('%Y%m%d')
                continue
            return date
        i += 1
    return None


def ASQuery_financial(code=None, report_date=None, fields=None, convert_column_name=False):
    """
    财务报表数据通用查询接口
    :param code: str/list
    :param report_date: str/list
    :param fields: list 财务字段英文名，参考ASData/financial_mean.py
    :param convert_column_name: bool 是否转换列名为中文，默认否
    :return: DataFrame
    """
    if code is None and report_date is None:
        raise ValueError('Specify at least one of code and report_date')

    filter_ = {}
    if code:
        if isinstance(code, str):
            filter_.__setitem__('code', code)
        elif isinstance(code, list):
            filter_.__setitem__('code', {'$in': code})
        else:
            raise TypeError('code must be str or list type')

    if report_date:
        if isinstance(report_date, str):
            filter_.__setitem__('report_date', int(report_date))
        elif isinstance(report_date, list):
            filter_.__setitem__('report_date', {'$in': [int(r) for r in report_date]})
        else:
            raise TypeError('report_date must be str or list type')

    projection = {'_id': 0}
    if fields:
        if not isinstance(fields, list):
            raise TypeError('fields must be list type')

        for f in fields:
            field_id = financial_dict_english_to_id.get(f, f)
            projection.__setitem__(field_id, 1)
        # if 'code' not in fields:
        #    projection.__setitem__('code', 1)
        # if 'report_date' not in fields:
        #     projection.__setitem__('report_date', 1)

    coll = settings.database.financial
    try:
        cursor = coll.find(filter_, projection)
    except Exception as e:
        print(e)
        return None

    data = [row for row in cursor]
    df = pd.DataFrame(data)
    if convert_column_name:
        df = df.rename(columns=lambda x: financial_dict_id_to_chinese.get(x, x))
    else:
        df = df.rename(columns=lambda x: financial_dict_id_to_english.get(x, x))
    return df


def ASQuery_stock_pool():
    coll = settings.database.stock_pool
    filter = {}
    projection = {'_id': 0}
    res = coll.find(filter, projection)
    res = [i for i in res]
    df = pd.DataFrame(res)
    return df
